"""
WebSocket API路由
"""
import asyncio
import json
from typing import Dict, Set
from datetime import datetime
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from fastapi.websockets import WebSocketState

from app.core.kafka_client import get_kafka_client
from app.utils.logger import get_logger

router = APIRouter()
logger = get_logger(__name__)


class WebSocketManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        # 存储活跃连接
        self.active_connections: Set[WebSocket] = set()
        # 存储Topic订阅
        self.topic_subscriptions: Dict[str, Set[WebSocket]] = {}
        # 存储消费任务
        self.consume_tasks: Dict[str, asyncio.Task] = {}
        # 存储监控订阅
        self.monitor_subscriptions: Set[WebSocket] = set()
        # 监控数据推送任务
        self.monitor_task: asyncio.Task = None
    
    async def connect(self, websocket: WebSocket):
        """接受WebSocket连接"""
        await websocket.accept()
        self.active_connections.add(websocket)
        logger.info(f"WebSocket连接建立，当前连接数: {len(self.active_connections)}")
    
    def disconnect(self, websocket: WebSocket):
        """断开WebSocket连接"""
        self.active_connections.discard(websocket)
        
        # 从所有Topic订阅中移除
        for topic, subscribers in self.topic_subscriptions.items():
            subscribers.discard(websocket)
            
            # 如果没有订阅者了，停止消费任务
            if not subscribers and topic in self.consume_tasks:
                self.consume_tasks[topic].cancel()
                del self.consume_tasks[topic]
        
        # 从监控订阅中移除
        self.monitor_subscriptions.discard(websocket)

        # 如果没有监控订阅者了，停止监控任务
        if len(self.monitor_subscriptions) == 0 and self.monitor_task:
            self.monitor_task.cancel()
            self.monitor_task = None

        logger.info(f"WebSocket连接断开，当前连接数: {len(self.active_connections)}")
    
    async def subscribe_topic(self, websocket: WebSocket, topic: str):
        """订阅Topic"""
        if topic not in self.topic_subscriptions:
            self.topic_subscriptions[topic] = set()
        
        self.topic_subscriptions[topic].add(websocket)
        
        # 如果是第一个订阅者，启动消费任务
        if len(self.topic_subscriptions[topic]) == 1:
            task = asyncio.create_task(self._consume_topic(topic))
            self.consume_tasks[topic] = task
        
        logger.info(f"WebSocket订阅Topic: {topic}")
    
    async def unsubscribe_topic(self, websocket: WebSocket, topic: str):
        """取消订阅Topic"""
        if topic in self.topic_subscriptions:
            self.topic_subscriptions[topic].discard(websocket)
            
            # 如果没有订阅者了，停止消费任务
            if not self.topic_subscriptions[topic] and topic in self.consume_tasks:
                self.consume_tasks[topic].cancel()
                del self.consume_tasks[topic]
                del self.topic_subscriptions[topic]
        
        logger.info(f"WebSocket取消订阅Topic: {topic}")
    
    async def _consume_topic(self, topic: str):
        """消费Topic消息并推送给订阅者"""
        try:
            kafka_client = get_kafka_client()

            # 创建专门用于实时流式传输的消费者
            consumer = await kafka_client.get_consumer(
                topic,
                from_beginning=False,  # 从最新位置开始
                unique_group=True      # 使用唯一组避免冲突
            )

            logger.info(f"开始实时消费Topic: {topic}")

            while True:
                # 获取当前Topic的订阅者
                subscribers = self.topic_subscriptions.get(topic, set())

                if not subscribers:
                    logger.info(f"Topic {topic} 没有订阅者，停止消费")
                    break  # 没有订阅者，退出

                try:
                    # 使用更短的超时时间和更小的批次大小来实现实时处理
                    msg_batch = await asyncio.wait_for(
                        consumer.getmany(timeout_ms=100, max_records=1),  # 100ms超时，每次最多1条消息
                        timeout=0.2  # 200ms总超时
                    )

                    # 处理批次中的消息
                    for tp, messages in msg_batch.items():
                        for msg in messages:
                            try:
                                # 处理消息值，支持各种类型
                                value = msg.value
                                if value:
                                    # 如果是字节类型，先解码
                                    if isinstance(value, bytes):
                                        try:
                                            value = value.decode('utf-8')
                                        except UnicodeDecodeError:
                                            # 如果解码失败，转换为base64
                                            import base64
                                            value = base64.b64encode(value).decode('ascii')
                                            value = f"[BASE64] {value}"

                                    # 尝试解析JSON
                                    if isinstance(value, str):
                                        try:
                                            import orjson
                                            value = orjson.loads(value)
                                        except:
                                            pass  # 保持原始字符串

                                # 处理消息键
                                key = msg.key
                                if key and isinstance(key, bytes):
                                    try:
                                        key = key.decode('utf-8')
                                    except UnicodeDecodeError:
                                        import base64
                                        key = base64.b64encode(key).decode('ascii')
                                        key = f"[BASE64] {key}"

                                # 准备消息数据
                                message_data = {
                                    "type": "message",
                                    "topic": topic,
                                    "data": {
                                        "topic": msg.topic,
                                        "partition": msg.partition,
                                        "offset": msg.offset,
                                        "timestamp": msg.timestamp,
                                        "key": key,
                                        "value": value,
                                        "headers": dict(msg.headers) if msg.headers else {}
                                    }
                                }

                                # 立即发送给所有订阅者
                                disconnected = set()
                                for websocket in subscribers:
                                    try:
                                        if websocket.client_state == WebSocketState.CONNECTED:
                                            # 立即发送消息，不等待
                                            await websocket.send_text(json.dumps(message_data))
                                            # 强制刷新WebSocket缓冲区
                                            await asyncio.sleep(0)  # 让出控制权，确保消息立即发送
                                        else:
                                            disconnected.add(websocket)
                                    except Exception as e:
                                        logger.error(f"发送WebSocket消息失败: {e}")
                                        disconnected.add(websocket)

                                # 清理断开的连接
                                for ws in disconnected:
                                    self.disconnect(ws)

                                # 记录消息发送日志（可选，用于调试）
                                logger.debug(f"实时发送消息: {topic} partition={msg.partition} offset={msg.offset}")

                            except Exception as e:
                                logger.error(f"处理消息时出错: {e}")
                                continue

                except asyncio.TimeoutError:
                    # 超时是正常的，继续循环等待新消息
                    continue
                except Exception as e:
                    logger.error(f"获取消息批次时出错: {e}")
                    await asyncio.sleep(1)  # 出错时稍等再试
                    continue

        except asyncio.CancelledError:
            logger.info(f"Topic消费任务被取消: {topic}")
        except Exception as e:
            logger.error(f"Topic消费任务出错: {e}")
        finally:
            # 确保消费者被正确关闭
            try:
                if 'consumer' in locals():
                    await consumer.stop()
            except Exception as e:
                logger.error(f"关闭消费者时出错: {e}")
    
    async def send_personal_message(self, websocket: WebSocket, message: dict):
        """发送个人消息"""
        try:
            if websocket.client_state == WebSocketState.CONNECTED:
                await websocket.send_text(json.dumps(message))
        except Exception as e:
            logger.error(f"发送个人消息失败: {e}")

    async def broadcast_json(self, message: dict):
        """向所有连接的客户端广播JSON消息"""
        if not self.active_connections:
            return

        disconnected_clients = set()
        for websocket in self.active_connections.copy():
            try:
                if websocket.client_state == WebSocketState.CONNECTED:
                    await websocket.send_text(json.dumps(message))
                else:
                    disconnected_clients.add(websocket)
            except Exception as e:
                logger.error(f"广播消息失败: {e}")
                disconnected_clients.add(websocket)

        # 清理断开的连接
        for ws in disconnected_clients:
            self.disconnect(ws)

    async def subscribe_monitor(self, websocket: WebSocket):
        """订阅监控数据"""
        self.monitor_subscriptions.add(websocket)

        # 如果还没有监控任务，启动一个
        if self.monitor_task is None or self.monitor_task.done():
            self.monitor_task = asyncio.create_task(self._push_monitor_data())

        logger.info(f"WebSocket订阅监控数据，当前监控订阅数: {len(self.monitor_subscriptions)}")

    async def unsubscribe_monitor(self, websocket: WebSocket):
        """取消订阅监控数据"""
        self.monitor_subscriptions.discard(websocket)

        # 如果没有订阅者了，停止监控任务
        if len(self.monitor_subscriptions) == 0 and self.monitor_task:
            self.monitor_task.cancel()
            self.monitor_task = None

        logger.info(f"WebSocket取消订阅监控数据，当前监控订阅数: {len(self.monitor_subscriptions)}")

    async def _push_monitor_data(self):
        """推送监控数据给订阅者"""
        try:
            while len(self.monitor_subscriptions) > 0:
                try:
                    # 获取监控数据
                    monitor_data = await self._get_monitor_data()

                    # 推送给所有监控订阅者
                    disconnected_clients = set()
                    for websocket in self.monitor_subscriptions.copy():
                        try:
                            if websocket.client_state == WebSocketState.CONNECTED:
                                # 将Pydantic模型转换为字典，然后序列化
                                serializable_data = {}

                                # 处理概览数据
                                if monitor_data.get("overview"):
                                    overview = monitor_data["overview"]
                                    if hasattr(overview, 'model_dump'):
                                        # Pydantic v2
                                        serializable_data["overview"] = overview.model_dump()
                                    elif hasattr(overview, 'dict'):
                                        # Pydantic v1
                                        serializable_data["overview"] = overview.dict()
                                    else:
                                        # 普通字典
                                        serializable_data["overview"] = overview

                                # 处理Topic统计数据
                                if monitor_data.get("topics"):
                                    topics = monitor_data["topics"]
                                    if isinstance(topics, list):
                                        serializable_topics = []
                                        for topic in topics:
                                            if hasattr(topic, 'model_dump'):
                                                serializable_topics.append(topic.model_dump())
                                            elif hasattr(topic, 'dict'):
                                                serializable_topics.append(topic.dict())
                                            else:
                                                serializable_topics.append(topic)
                                        serializable_data["topics"] = serializable_topics
                                    else:
                                        serializable_data["topics"] = topics

                                # 处理性能指标数据
                                if monitor_data.get("performance"):
                                    performance = monitor_data["performance"]
                                    if hasattr(performance, 'model_dump'):
                                        serializable_data["performance"] = performance.model_dump()
                                    elif hasattr(performance, 'dict'):
                                        serializable_data["performance"] = performance.dict()
                                    else:
                                        serializable_data["performance"] = performance

                                message_data = {
                                    "type": "monitor_update",
                                    "data": serializable_data,
                                    "timestamp": datetime.now().isoformat()
                                }

                                # 添加调试日志
                                logger.debug(f"推送监控数据: overview={bool(serializable_data.get('overview'))}, "
                                           f"topics={len(serializable_data.get('topics', []))}, "
                                           f"performance={bool(serializable_data.get('performance'))}")

                                await websocket.send_text(json.dumps(message_data, default=str))
                            else:
                                disconnected_clients.add(websocket)
                        except Exception as e:
                            logger.warning(f"推送监控数据失败: {e}")
                            disconnected_clients.add(websocket)

                    # 清理断开的连接
                    for websocket in disconnected_clients:
                        self.monitor_subscriptions.discard(websocket)

                    # 等待10秒后再次推送（更频繁的实时更新）
                    await asyncio.sleep(10)

                except Exception as e:
                    logger.error(f"获取监控数据失败: {e}")
                    await asyncio.sleep(10)  # 出错时等待较短时间

        except asyncio.CancelledError:
            logger.info("监控数据推送任务被取消")
        except Exception as e:
            logger.error(f"监控数据推送任务出错: {e}")

    async def _get_monitor_data(self):
        """获取监控数据"""
        try:
            # 这里调用监控API的逻辑来获取数据
            # 为了避免循环导入，我们直接在这里实现简化的监控数据获取
            from app.api.monitor import get_monitor_overview, get_topic_statistics, get_performance_metrics

            # 分别获取监控数据，即使某些失败也继续
            monitor_data = {}

            # 获取概览数据
            try:
                overview = await get_monitor_overview()
                monitor_data["overview"] = overview
            except Exception as e:
                logger.warning(f"获取概览数据失败: {e}")
                monitor_data["overview"] = None

            # 获取Topic统计
            try:
                topics = await get_topic_statistics()
                monitor_data["topics"] = topics
            except Exception as e:
                logger.warning(f"获取Topic统计失败: {e}")
                monitor_data["topics"] = []

            # 获取性能指标
            try:
                performance = await get_performance_metrics()
                monitor_data["performance"] = performance
            except Exception as e:
                logger.warning(f"获取性能指标失败: {e}")
                monitor_data["performance"] = None

            return monitor_data

        except Exception as e:
            logger.error(f"获取监控数据失败: {e}")
            return {
                "overview": None,
                "topics": [],
                "performance": None,
                "error": str(e)
            }


# 全局WebSocket管理器
manager = WebSocketManager()


@router.websocket("/stream")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket流式数据端点"""
    await manager.connect(websocket)
    
    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            
            try:
                message = json.loads(data)
                message_type = message.get("type")
                
                if message_type == "subscribe":
                    # 订阅Topic
                    topic = message.get("topic")
                    if topic:
                        await manager.subscribe_topic(websocket, topic)
                        await manager.send_personal_message(websocket, {
                            "type": "subscribed",
                            "topic": topic,
                            "message": f"已订阅Topic: {topic}"
                        })
                
                elif message_type == "unsubscribe":
                    # 取消订阅Topic
                    topic = message.get("topic")
                    if topic:
                        await manager.unsubscribe_topic(websocket, topic)
                        await manager.send_personal_message(websocket, {
                            "type": "unsubscribed",
                            "topic": topic,
                            "message": f"已取消订阅Topic: {topic}"
                        })
                
                elif message_type == "subscribe_monitor":
                    # 订阅监控数据
                    await manager.subscribe_monitor(websocket)
                    await manager.send_personal_message(websocket, {
                        "type": "monitor_subscribed",
                        "message": "已订阅监控数据"
                    })

                elif message_type == "unsubscribe_monitor":
                    # 取消订阅监控数据
                    await manager.unsubscribe_monitor(websocket)
                    await manager.send_personal_message(websocket, {
                        "type": "monitor_unsubscribed",
                        "message": "已取消订阅监控数据"
                    })

                elif message_type == "ping":
                    # 心跳检测
                    await manager.send_personal_message(websocket, {
                        "type": "pong",
                        "timestamp": message.get("timestamp")
                    })

                else:
                    await manager.send_personal_message(websocket, {
                        "type": "error",
                        "message": f"未知消息类型: {message_type}"
                    })
                    
            except json.JSONDecodeError:
                await manager.send_personal_message(websocket, {
                    "type": "error",
                    "message": "无效的JSON格式"
                })
            except Exception as e:
                logger.error(f"处理WebSocket消息失败: {e}")
                await manager.send_personal_message(websocket, {
                    "type": "error",
                    "message": f"处理消息失败: {str(e)}"
                })
                
    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        logger.error(f"WebSocket连接出错: {e}")
        manager.disconnect(websocket)
